import os

import dotenv
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent

dotenv.load_dotenv()

# 1 构建LLM
llm = ChatOpenAI(
    model="glm-4",
    base_url=os.getenv("ZP_BASE"),
    api_key=os.getenv("ZP_API_KEY"),
    temperature=0
)

# 2 构建agent
agent = create_react_agent(
    model=llm,
    tools=[],
    prompt="你是一个智能助手"
)

# print(agent.invoke({"messages": [{"role": "user", "content": "你是谁?能帮我解决什么问题?"}]}))

# 3 流式输出
# updates: 流式输出每个工具调用的每个步骤
# messages: 流式输出大语言模型回复的Token
# values: 一次拿到所有的chunk,默认值
# custom: 自定义输出.主要是可以在工具内部使用get_stream_writer获取输入流,添加自定义内容
for chunk in agent.stream(
        {"messages": [{"role": "user", "content": "你是谁?能帮我解决什么问题?"}]},
        # stream_mode="messages",
):
    print(chunk, flush=True)
